import time
from ddt import ddt, data
from openmind_hub import create_repo, delete_repo, model_info, list_models, get_model_ci_info
import unittest
from data import license_data
from openmind_hub import RepositoryNotFoundError, GatedRepoError
from openmind_hub.plugins.openmind.utils._error import BadRequestError


repo_id = "guoxiaozhen2/model_test2024"
license = "mit"
token = '123456789()'



@ddt
class TestModel(unittest.TestCase):

    def test_01_create_modelrepo_success(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                # token= get_config().get("'DEFAULT'", "'token'"),
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            a = model_info(
                repo_id=repo_id,
                # token= get_config().get("'DEFAULT'", "'token'"),
                token=token,
                timeout=None
            )
            self.assertIsNotNone(a, "仓库不存在，创建失败")
        except:
            raise Exception("创建仓库有报错")
        else:
            delete_repo(
                repo_id=repo_id,
                # token= get_config().get("'DEFAULT'", "'token'"),
                token=token,
                repo_type="model",
                missing_ok=True
            )

    def test_02_create_modelrepo_duplicate_existok(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
        except:
            raise Exception("exist_ok=True未生效")
        else:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )

    def test_03_create_modelrepo_duplicate_existnotok(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=False,
                fullname=None
            )
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=False,
                fullname=None
            )
            raise Exception("exist_ok=False未生效")
        except BadRequestError:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )

    def test_04_delete_modelrepo_success(self):
        author = "guoxiaozhen2"
        try:
            a = list_models(
                filter=None,
                author=author,
                token=token
            )
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=False,
                fullname=None
            )
            time.sleep(3)
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
            b = list_models(
                filter=None,
                author=author,
                token=token
            )
            self.assertEqual(list(a), list(b), "仓库存在，删除失败")
        except:
            raise Exception("删除仓库有报错")

    def test_05_delete_modelrepo_duplicate_missok(self):
        try:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
        except:
            raise Exception("missing_ok=True未生效")

    def test_06_delete_modelrepo_duplicate_missnotok(self):
        try:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=False
            )
            raise Exception("missing_ok=False未生效")
        except RepositoryNotFoundError:
            pass

    def test_07_create_modelrepo_wrong_token(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token + "a",
                repo_type="model",
                exist_ok=False,
                fullname=None
            )
            raise Exception("token未校验")
        except GatedRepoError:
            pass

    def test_08_create_modelrepo_no_repo_type(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=True,
                desc="描述内容",
                license=license,
                token=token,
                exist_ok=False,
                fullname=None
            )
        except:
            raise Exception("仓库类型未添加默认值")
        else:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=False
            )

    def test_09_create_modelrepo_wrong_license(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license + "2",
                token=token,
                repo_type="model",
                exist_ok=False,
                fullname=None
            )
            raise Exception("license未校验")
        except BadRequestError:
            pass

    @data(*license_data.values())
    def test_10_create_modelrepo_morelicenses(self, data):
        license = data["license"]
        print(license)
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            a = model_info(
                repo_id=repo_id,
                token=token,
                timeout=None
            )
            self.assertIsNotNone(a, "仓库不存在，创建失败")
        except:
            raise Exception("创建仓库有报错")
        else:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
    def test_11_get_model_ci_info_success(self):
        repo_id="CITesting/chatglm2_6b"
        try:
            a = get_model_ci_info(
                repo_id=repo_id,
                timeout=None,
                token=token,
            )
            self.assertEqual(a.total, 1)
        except:
            raise Exception("获取模型CI信息有误")
    def test_12_get_model_ci_info_wrong_token(self):
        repo_id="CITesting/chatglm2_6"
        try:
            a = get_model_ci_info(
                repo_id=repo_id,
                timeout=None,
                token=token+"a",
            )
            raise Exception("token未校验")
        except GatedRepoError:
            pass
    def test_13_get_model_ci_info_noexistrepo(self):
        repo_id="CITesting/chatglm2_6"
        try:
            a = get_model_ci_info(
                repo_id=repo_id+"a",
                timeout=None,
                token=token,
            )
            raise Exception("repo未校验")
        except RepositoryNotFoundError:
            pass



if __name__ == '__main__':
    unittest.main()
